from langchain_core.prompts import ChatPromptTemplate
from pydantic import BaseModel, Field
from langchain_core.output_parsers import JsonOutputParser


class Outline(BaseModel):
    title: str = Field(description="一级目录")
    sub_titles: list[str] = Field(description="二级目录列表")


class Scheme(BaseModel):
    abstract: str = Field(description="全文的摘要")
    outlines: list[Outline] = Field(description="项目大纲")


_planner_system_template = """
你是一个优秀的方案规划师，
你的职责是根据标题，规划方案大纲，
要求:
1.大纲至少要规划两个层级以上，
2.大纲规划的层级中不需要序号，
3.另外还需要编写一个全文摘要，更好的为后续生成方案服务

输出：
{output_format}
"""

_planner_human_template = """
题目：{project_name}
"""


class Planner:

    def __init__(self, llm):

        self._llm = llm
        _prompt = ChatPromptTemplate.from_messages(
            [
                ("system", _planner_system_template),
                ("human", _planner_human_template),
            ]
        )
        _parser = JsonOutputParser(pydantic_object=Scheme)
        _prompt = _prompt.partial(output_format=_parser.get_format_instructions())
        self._chain = _prompt | self._llm | _parser

    def __call__(self, state):
        return self._chain.invoke(state)


if __name__ == "__main__":
    from langchain_openai import ChatOpenAI

    _llm = ChatOpenAI(
        api_key="ollama",
        model="qwen2.5:7b",
        base_url="http://192.168.10.13:60001/v1",
        temperature=0.7,
    )
    planner = Planner(_llm)
    result = planner({"project_name": "基于多模态大模型驱动的数字人"})
    print(result["abstract"])

